import numpy as np
import matplotlib.pyplot as plt


# 生成一个示例二维图像
def generate_image(size=128):
    """
    生成一个简单的二维图像，包含几个亮点。

    :param size: 图像的大小（size x size）
    :return: 生成的二维图像
    """
    image = np.zeros((size, size), dtype=np.float64)

    # 在图像中添加一些亮点
    image[size // 4, size // 4] = 255
    # image[3 * size // 4, size // 4] = 255
    # image[size // 2, size // 2] = 255
    image[size // 4, 3 * size // 4] = 255
    # image[3 * size // 4, 3 * size // 4] = 255

    return image


# 生成原始图像
size = 128
original_image = generate_image(size)

# 进行二维傅里叶变换
fft_image = np.fft.fft2(original_image)
fft_image_shifted = np.fft.fftshift(fft_image)  # 将零频率成分移动到中心

# 计算傅里叶变换的幅度谱
magnitude_spectrum = np.log(np.abs(fft_image_shifted) + 1)  # 取对数以便于可视化


# 进行二维傅里叶反变换
ifft_image = np.fft.ifft2(fft_image).real  # 取实部

# 可视化三个图像
plt.figure(figsize=(15, 5))

# 原始图像
plt.subplot(1, 3, 1)
plt.imshow(original_image)
plt.title('Original Image')
plt.colorbar()

# 傅里叶变换的幅度谱
plt.subplot(1, 3, 2)
plt.imshow(magnitude_spectrum)
plt.title('Magnitude Spectrum (FFT)')
plt.colorbar()

# 反变换后的图像
plt.subplot(1, 3, 3)
plt.imshow(ifft_image)
plt.title('Reconstructed Image (IFFT)')
plt.colorbar()

plt.tight_layout()
plt.show()